#include <stdio.h>
#include <yolo.h>
#include <opencv2/opencv.hpp>

int main(int argc, char **argv)
{
    // 检查输入参数
    if(argc != 3){
        printf("Usage: %s <model_path> <image_path>\n", argv[0]);
        return -1;
    }

    // 初始网络模型
    char *model_path = argv[1]; // 模型加载路径
    YOLO model = YOLO(model_path, 640, 640, 80);

    // 读取检测图片
    char *image_path = argv[2];             // 读取图片路径
    cv::Mat image = cv::imread(image_path); // 读取图片文件
    if(image.empty()){
        printf("Input image is empty!\n");
        return -1;
    }

    // 进行目标检测
    int det_num = 0;             // 检测目标数量
    std::vector<Result> results; // 检测结果向量

    struct timeval tv1, tv2;
	long t1, t2, time;
    gettimeofday(&tv1, NULL);

    det_num = model.detect(image, results, 0.25f, 0.45f);

    gettimeofday(&tv2, NULL);
    t1 = tv2.tv_sec - tv1.tv_sec;
    t2 = tv2.tv_usec - tv1.tv_usec;
    time = (long)(t1 * 1000 + t2 / 1000);
    printf("Detect times: %dms\n", time);

    // 保存结果图像
	if(det_num > 0){
		// 绘制检测结果
		for(int i = 0; i < det_num; i++){
			// 绘制目标边框
			cv::Rect bbox = cv::Rect(results[i].box_x, results[i].box_y, results[i].box_w, results[i].box_h);
			cv::rectangle(image, bbox, cv::Scalar(0, 255, 0), 2);

			// 绘制目标信息
			char text[256];
			sprintf(text, "label: %d score: %0.2f", results[i].label, results[i].score);
			cv::putText(image, text, bbox.tl(), cv::FONT_HERSHEY_COMPLEX, 1, cv::Scalar(0, 0, 255), 1);
		}

		// 保存结果图像
		cv::imwrite("result.jpg", image);
		printf("Writing result.jpg success!\n");
	}

    return 0;
}